import os

from catalog.place.place_api import PlaceAPI
from catalog.models.sentinel2_l2a import Sentinel2L2AScene
from catalog.raster.raster import retrieve_by_uid
from common.geoimage.raster_dataset import mosaic_raster_files, mosaic_raster_datasets
from common.config import DEFAULT_CACHE_DIR
from common.io_utils.dirs import mkdir_p

aoi_list = PlaceAPI.get_geom_list_with_place_name(
    name_en="Morgan",
    region_level=3,
    parent_name_en="United States:Illinois:Morgan",
    name_exact_mode=True,
)
aoi_geom = aoi_list[0]
aoi_geom.to_shapely()

items = Sentinel2L2AScene.query_many_items(
    spatial=aoi_geom, temporal=["2021-01-04", "2021-01-8"]
)
# for item in items:
#    print(item.uid, item.epsg_no)

rdss = []
dir_data_tmp = os.path.join(DEFAULT_CACHE_DIR, "test_mosaic")
mkdir_p(dir_data_tmp)
fpaths_data_tmp = []
mkdir_p(dir_data_tmp)
for item in items:
    rds = retrieve_by_uid(
        Sentinel2L2AScene.collection_id,
        item.uid,
        asset_names=["B04", "B03", "B02"],
        dest_res=10,
    )
    print(rds.meta)
    rdss.append(rds)
    fpath_data = os.path.join(dir_data_tmp, f"{item.uid}.tif")
    rds.to_geotiff(fpath_data)
    fpaths_data_tmp.append(fpath_data)

print(fpaths_data_tmp)
fpath_dest = os.path.join(dir_data_tmp, "mosaic.tif")

mosaicked_rds = mosaic_raster_files(
    fpaths_data_tmp,
    resample_alg="bilinear",
    dest_aoi=aoi_geom,
    crop_to_aoi=True,
)

mosaicked_rds = mosaic_raster_datasets(
    rdss,
    resample_alg="bilinear",
    dest_aoi=aoi_geom,
    crop_to_aoi=True,
)
print(mosaicked_rds.meta)
